'''
* This is the projet for Brtc LlmOps Platform
* @Author Leon-liao <liaosiliang@alltman.com>
* @Description //TODO 
* @File: 2_study_rrf_strategy.py
* @Time: 2025/10/30
* @All Rights Reserve By Brtc
'''
import dotenv
import weaviate
from langchain.retrievers import MultiQueryRetriever
from langchain_core.callbacks import CallbackManagerForRetrieverRun
from langchain_core.documents import Document
from langchain_core.load import dumps, loads
from langchain_openai import OpenAIEmbeddings, ChatOpenAI
from langchain_weaviate import WeaviateVectorStore

dotenv.load_dotenv()
class RAGFusionRetriever(MultiQueryRetriever):
    """RAG多查询结果融合检索器"""
    k:int = 4

    def __init__(self, k:int=4, **kwargs):
        super().__init__(**kwargs)
        self.k = k

    def retrieve_documents( self, queries: list[str], run_manager: CallbackManagerForRetrieverRun, ) -> list[Document]:
        """重写检索文档，返回二层文档的嵌套列表"""
        documents = []
        #将query 分解成 3个问题
        for query in queries:
            # 每个问题都进行检索出4个相关文档
            docs = self.retriever.invoke(
                query, config={"callback": run_manager.get_child()}
            )
            # 再把文档追加到列表中
            documents.append(docs)
        #3*4， 3*4 维度的列表
        return documents


    def unique_union(self, documents: list[Document]) -> list[Document]:
        """使用RRF算法对文档列表进行排序&合并"""
        #1、初始化一个字典,用于存储每一个唯一的文档得分
        fused_score = {}
        #2、遍历每个查询对应的文档列表
        for  docs in documents:
            #3、内层遍历文档列表得得到每一个文档的得分
            for  rank, doc in enumerate(docs):
                #4、将文档使用langchain提供的工具转换成字符串
                doc_str = dumps(doc)
                #5、检测该字符串是否存在得分， 如不不存在则赋值为0 分
                if doc_str not in fused_score:
                    fused_score[doc_str] = 0
                #6、计算多结果得分， 得分越高越考前，k 为权重参数
                fused_score[doc_str]  += 1/(rank+60)
        #7、提取得分并进行排序
        rerank_results = [
            (loads(doc), score)
            for doc, score in sorted(fused_score.items(), key=lambda x:x[1], reverse=True)
        ]

        return [item[0] for item in rerank_results[:self.k]]
client = weaviate.connect_to_local("192.168.106.129", 8080)
db = WeaviateVectorStore(
    client,
    index_name="TestDemo",
    text_key="text",
    embedding=OpenAIEmbeddings(model="text-embedding-3-small")
)
retriever = db.as_retriever(search_type="mmr")
rag_fusion_retriever = RAGFusionRetriever.from_llm(
    retriever=retriever,
    llm=ChatOpenAI(model="gpt-4o-mini", temperature=0.0)
)

# 执行检索
docs = rag_fusion_retriever.invoke("关于llmops 的应用配置文档有哪些？")
for doc in docs:
    print("=========================================")
    print(doc.page_content[:50])
client.close()
